import logging
import numpy as np
from scipy.misc import imread
from scipy.misc import imsave
from scipy.signal import convolve2d
from math import floor
import math


def getFilter(name, size, cutoffFreq):
    # assert size[0] % 2 == 1
    # assert size[1] % 2 == 1
    assert cutoffFreq > 0
    crow = floor(size[0] / 2)
    ccol = floor(size[1] / 2)
    dist = np.zeros(size)
    for row in range(size[0]):
        for col in range(size[1]):
            dist[row, col] = (row - crow) ** 2 + (col - ccol) ** 2
    # print(dist)
    if name == 'gaussian':
        return np.exp(-dist / 2.0 / (cutoffFreq ** 2))
    elif name == 'idea':
        window = (dist <= cutoffFreq ** 2) + 0
        return window
    elif name == 'butterWorth':
        return 1 / ((dist / (cutoffFreq ** 2) + 1))**2
    else:
        logging.log(logging.ERROR, "No such filter: " + name)

def ideal(matrix, cutoff, passband):
    newData = matrix.copy()
    # integer division in python is a suprise
    center = (math.floor(newData.shape[0]/2), math.floor(newData.shape[1]/2))

    for (x, y), value in np.ndenumerate(matrix):
        dist = math.sqrt((x-center[0])**2+(y-center[1])**2)
        if passband == 'low':
            # leave off high frequency
            if dist > cutoff:
                newData[x][y] = 0 + 0j
        if passband == 'high':
            # cut off low frequency
             if dist < cutoff:
                newData[x][y] = 0 + 0j
    return newData


def main():
    FORMAT = '	[%(levelname)-5s]%(asctime)-8s %(filename)s:%(lineno)d %(message)s'
    DATEFORMAT = '%H:%M:%S'
    logging.basicConfig(level=logging.DEBUG, format=FORMAT, datefmt=DATEFORMAT)
    logging.log(logging.DEBUG, "Start")

    oriimg = imread("/home/d/Documents/lu-ht/DIP/images/characters_test_pattern.tif", mode='F')
    oriimg = oriimg
    print(oriimg)
    CutOffFreq = [5] #, 15, 30]
    FilterName = ['idea'] #, 'butterWorth', 'gaussian']
    zeroPadding = np.zeros(oriimg.shape)
    for row in range(zeroPadding.shape[0]):
        for col in range(zeroPadding.shape[1]):
            if (col + row) % 2:
                zeroPadding[row,col] = -1
            else:
                zeroPadding[row,col] = 1

    logging.log(logging.DEBUG, zeroPadding)

    for method in FilterName:
        for freq in CutOffFreq:
            window = getFilter(method, oriimg.shape, freq)
            preprocess = np.multiply(oriimg, zeroPadding)
            fourierTransform = np.fft.fft2(preprocess)
            filtered_img = np.multiply(fourierTransform, window)
            # filtered_img = ideal(fourierTransform, freq, 'low')
            inverseFourier = np.fft.ifft2(filtered_img)
            inverseFourier = np.real(inverseFourier)
            outImg = np.multiply(inverseFourier, zeroPadding)
            # outImg[outImg > 255] = 255
            imsave("%s_Lowpass_%d.jpg" % (method, freq), outImg)
            logging.log(logging.DEBUG, '%f %f' % (np.max((outImg)), np.min(outImg)))

    CutOffFreq = [15, 30, 80]
    for method in FilterName:
        for freq in CutOffFreq:
            window = 1 - getFilter(method, oriimg.shape, freq)
            logging.log(logging.DEBUG, "Max %f, min %f" % (np.max(window), np.min(window)))
            # preprocess = oriimg
            preprocess = np.multiply(oriimg, zeroPadding)
            fourierTransform = np.fft.fft2(preprocess)
            filtered_img = np.multiply(fourierTransform,window)
            inverseFourier = np.fft.ifft2(filtered_img)
            inverseFourier = np.real(inverseFourier)
            # outImg = inverseFourier
            outImg = np.multiply(inverseFourier, zeroPadding)
            outImg[outImg < 0] = 0
            outImg[outImg > 255] = 255
            imsave("%s_High_%d.jpg" % (method, freq), outImg)
            imsave("%s_High_filter_%d.jpg" % (method, freq), window)
    return 0


if __name__ == '__main__':
    main()
